#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Piper机械臂Xbox手柄控制系统 V4.0
作者: Yinchuan Li
日期: 2025-08-07

V4.0 版本升级说明:
1. 支持一个Xbox手柄同时控制两台Piper机械臂（can0、can1）。
2. 按下A键（按钮0）：切换为控制can0机械臂。
3. 按下B键（按钮1）：切换为控制can1机械臂。
4. 当前控制对象切换时，控制台有明显提示。
5. 其余控制逻辑与V3.0完全一致，所有操作均可无缝切换目标机械臂。
"""

# ...existing code...
import math
import os
import sys
import time

import mujoco
import numpy as np
import pygame
from mujoco import minimize
from piper_sdk import *
from scipy.spatial.transform import Rotation

# ...existing code...

# ...existing code...

class MujocoIKSolver:
    """
    The radius is the weight of the orientation residual (in Lie algebra)
    The smaller the radius, the more the orientation residual is ignored.

    Forward functions in mujoco:

    mujoco.mj_kinematics(self.model, self.data):
        compute the kinematics of the model and update the data inplace

    mujoco.mj_forward(self.model, self.data)
        = {
        1. mj_kinematics(model, data),  # update the position and the rotations
        2. compute the constraint forces and the accelerations
    }

    mujoco.mj_step(self.model, self.data):
        = {
        1. mj_forward(model, data),  # update the position and the rotations, forces and accelerations
        2. Do the integration of the dynamics
    }

    In this case, only position and rotation are updated, so mj_kinematics is enough.

    Mujoco use quaternion as [w, x, y, z]
    pinocchio use quaternion as [x, y, z, w]
    """

    def __init__(self, radius=0.04, reg=0.01):
        self.model = mujoco.MjModel.from_xml_path("piper_description/mujoco_model/piper_description.xml")
        self.data = mujoco.MjData(self.model)

        # Reset the state to the "zero_position" keyframe.
        self.go_zero()

        # Bounds at the joint limits.
        self.bounds = [self.model.jnt_range[:, 0], self.model.jnt_range[:, 1]]

        # Inital guess is the 'home' keyframe.
        self.j0 = self.model.key("home").qpos
        self.radius = radius
        self.reg = reg

    def go_zero(self):
        key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, "zero_position")
        mujoco.mj_resetDataKeyframe(self.model, self.data, key)
        mujoco.mj_forward(self.model, self.data)

    def move_target_to(self, pos, quat):
        self.data.mocap_pos[self.model.body("target").mocapid] = pos
        self.data.mocap_quat[self.model.body("target").mocapid] = quat
        mujoco.mj_kinematics(self.model, self.data)

    def ik_residual(self, joint, pos=None, quat=None, reg_target=None):
        """Residual for inverse kinematics.

        Args:
        joint: joint angles.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        radius: scaling of the 3D cross.

        Returns:
        The residual of the Inverse Kinematics task.
        """
        # Move the mocap body to the target
        tar_id = self.model.body("target").mocapid
        self.data.mocap_pos[tar_id] = self.model.body("target").pos if pos is None else pos
        self.data.mocap_quat[tar_id] = self.model.body("target").quat if quat is None else quat

        # Set qpos, compute forward kinematics.
        res = []
        for i in range(joint.shape[1]):
            self.data.qpos = joint[:, i]
            mujoco.mj_kinematics(self.model, self.data)

            # Position residual: can also use Lie algebra to compute the residual
            res_pos = self.data.site("effector").xpos - self.data.site("target").xpos

            # Effector quat, use mju_mat2quat.
            effector_quat = np.empty(4)
            mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)

            # Target quat, exploit the fact that the site is aligned with the body.
            target_quat = self.data.body("target").xquat

            # Orientation residual: quaternion difference.
            res_quat = np.empty(3)
            mujoco.mju_subQuat(res_quat, target_quat, effector_quat)
            res_quat *= self.radius

            # Regularization residual. To avoid the joint angles from being too far from the home position.
            # reg_target=x_prev is to regularize the joint angles to be close to the previous joint angles
            reg_target = self.model.key("home").qpos if reg_target is None else reg_target
            res_reg = self.reg * (joint[:, i] - reg_target)

            res_i = np.hstack((res_pos, res_quat, res_reg))
            res.append(np.atleast_2d(res_i).T)

        return np.hstack(res)

    def ik_jac(self, joint, res, pos=None, quat=None):
        """Analytic Jacobian of inverse kinematics residual

        Args:
        joint: joint angles.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        radius: scaling of the 3D cross.

        Returns:
        The Jacobian of the Inverse Kinematics task.
        """
        # least_squares() passes the value of the residual at x which is sometimes
        # useful, but we don't need it here.
        del res

        # Call mj_kinematics and mj_comPos (required for Jacobians).
        mujoco.mj_kinematics(self.model, self.data)
        mujoco.mj_comPos(self.model, self.data)

        # Get end-effector site Jacobian.
        jac_pos = np.empty((3, self.model.nv))
        jac_quat = np.empty((3, self.model.nv))
        mujoco.mj_jacSite(self.model, self.data, jac_pos, jac_quat, self.data.site("effector").id)

        # Get Deffector, the 3x3 mju_subquat Jacobian
        effector_quat = np.empty(4)
        mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)
        target_quat = self.data.body("target").xquat
        Deffector = np.empty((3, 3))
        # compute the jacobian: partial(target_quat - effector_quat) / partial(Deffector_quat)
        mujoco.mjd_subQuat(target_quat, effector_quat, None, Deffector)

        # Rotate into target frame, multiply by subQuat Jacobian, scale by radius.
        target_mat = self.data.site("target").xmat.reshape(3, 3)
        mat = self.radius * Deffector.T @ target_mat.T
        jac_quat = mat @ jac_quat

        # Regularization Jacobian.
        jac_reg = self.reg * np.eye(self.model.nv)

        return np.vstack((jac_pos, jac_quat, jac_reg))

    def solve_ik(self, pos=None, quat=None, j0=None, x_prev=None):
        """
        Solve the IK problem.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        j0: initial guess of the joint angles. If None, use the home position.
        """
        if j0 is None:
            # initial guess is the home position
            j0 = self.j0
        ik_target = lambda joint: self.ik_residual(joint, pos=pos, quat=quat, reg_target=x_prev)
        jac_target = lambda joint, r: self.ik_jac(joint, r, pos=pos, quat=quat)

        joint, _ = minimize.least_squares(
            j0, ik_target, self.bounds, jacobian=jac_target, verbose=0, check_derivatives=False, max_iter=1000
        )
        return joint

    def solve_ik_xyz_rpy(self, pos=None, rpy=None, j0=None, x_prev=None):
        """
        Solve the IK problem.
        pos: target position for the end effector.
        rpy: row, pitch, yaw of the end effector (in radians)
        j0: initial guess of the joint angles. If None, use the home position.
        """
        xyz_quat = Rotation.from_euler("xyz", rpy).as_quat()  # (x, y, z, w)
        quat = np.array([xyz_quat[3], xyz_quat[0], xyz_quat[1], xyz_quat[2]])  # (w, x, y, z)
        return self.solve_ik(pos=pos, quat=quat, j0=j0, x_prev=x_prev)


class XboxPiperControllerV3:
    """
    Xbox手柄控制Piper机械臂类 V3.0
    支持6自由度实时控制（位置+旋转）和夹爪操作
    优化控制精度和运动平滑性
    """

    def __init__(self, piper):
        """
        初始化Xbox控制器 V3.0

        Args:
            piper: Piper机械臂接口对象
        """
        self.piper = piper

        # 初始化pygame和手柄
        self._init_xbox_controller()

        # 机械臂控制参数 - V3.0超精细优化参数
        self.factor = 1000  # 单位转换因子（毫米转微米，度转微度）
        self.arm_speed = 50  # 机械臂速度（cm/s）
        self.move_speed = 10  # 移动速度（厘米/秒）- V3.0进一步降低为2cm/s，更加精细
        self.rotation_speed = 50  # 旋转速度（度/秒）- V3.0进一步降低为10°/s，更加平滑
        self.control_frequency = 120.0  # 控制频率（Hz）- V3.0保持60Hz
        self.move_distance_per_frame = self.move_speed / self.control_frequency  # 每帧移动距离（厘米）
        self.rotation_angle_per_frame = self.rotation_speed / self.control_frequency  # 每帧旋转角度（度）

        # 当前位置和初始位置
        self.current_position = None  # 当前机械臂位置 [X, Y, Z, RX, RY, RZ]
        self.initial_position = None  # 初始位置
        self.safe_initial_pose = [57.0, 0.0, 215.0, 0.0, 85.0, 0.0]  # 安全初始位姿

        # 摇杆控制参数 - V3.0超精细优化
        self.deadzone = 0.03  # 摇杆死区，V3.0进一步降低为0.03，提升微小控制响应性
        self.last_joystick_state = {
            "left_x": 0,
            "left_y": 0,
            "right_x": 0,
            "right_y": 0,
            "lt": 0,
            "rt": 0,
            "dpad_x": 0,
            "dpad_y": 0,
        }  # 记录上一次摇杆状态

        # 按钮状态记录
        self.last_button_states = {}
        self.last_hat_states = {}

        # 夹爪控制参数
        self.gripper_delay = 0.3  # 夹爪操作延时（秒）

        # V3.0新增：运动平滑性参数 - 超精细优化
        self.smoothness_factor = 0.9  # 运动平滑因子（0-1之间）- 提升到0.9，更加平滑
        self.last_target_position = None  # 上一次目标位置
        self.motion_smoothing_enabled = True  # 运动平滑开关
        self.velocity_smoothing = True  # 速度平滑开关
        self.last_velocity = [0.0] * 6  # 上一次速度
        self.velocity_smooth_factor = 0.7  # 速度平滑因子

        # V3.0新增：性能监控参数
        self.frame_count = 0
        self.last_fps_check = time.time()
        self.actual_fps = 0

        # 初始化机械臂
        self._initialize_robot()

        # add ik solver
        self.ik = MujocoIKSolver()

    def IKmove(self, X, Y, Z, RX, RY, RZ, speed=40):
        """
        X, Y, Z, RX, RY, RZ, J0 are machine values (int) of End effector
        """
        print(f"IKmoveJoint: X={X}, Y={Y}, Z={Z}, RX={RX}, RY={RY}, RZ={RZ}, speed={speed}")
        pos = np.array([X, Y, Z]) / 1e6
        rpy = np.deg2rad(np.array([RX, RY, RZ]) / 1000)
        # get current joints:
        arm_status = self.piper.GetArmJointMsgs().joint_state
        jo = [
            arm_status.joint_1,
            arm_status.joint_2,
            arm_status.joint_3,
            arm_status.joint_4,
            arm_status.joint_5,
            arm_status.joint_6,
            0,  # gripper 1
            0,
        ]
        j0 = np.deg2rad(np.asarray(jo) / 1000)
        j = self.ik.solve_ik_xyz_rpy(pos=pos, rpy=rpy, j0=j0)
        machine_joints = (np.rad2deg(j[:6]) * 1000).astype(int)
        self.piper.ModeCtrl(0x01, 0x01, speed, 0x00)
        self.piper.JointCtrl(*machine_joints)

    def _init_xbox_controller(self):
        """初始化Xbox手柄"""
        pygame.init()
        pygame.joystick.init()

        if pygame.joystick.get_count() == 0:
            raise RuntimeError("未检测到Xbox手柄，请确保手柄已连接")

        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()

        print(f"已连接Xbox手柄: {self.joystick.get_name()}")
        print(f"轴数量: {self.joystick.get_numaxes()}, 按钮数量: {self.joystick.get_numbuttons()}")
        print(f"帽子数量: {self.joystick.get_numhats()}")

        # Xbox按键映射
        self.button_mapping = {
            4: "LB",  # 左肩键 - 关闭夹爪
            5: "RB",  # 右肩键 - 打开夹爪
            7: "Menu",  # 菜单键 - 重置位置
        }

        # 摇杆轴映射
        self.axis_mapping = {
            0: "Left Stick X",  # 左摇杆X轴 - 控制Y轴移动
            1: "Left Stick Y",  # 左摇杆Y轴 - 控制X轴移动
            2: "LT Trigger",  # 左触发器 - Z轴向下
            3: "Right Stick X",  # 右摇杆X轴 - 控制RZ轴旋转
            4: "Right Stick Y",  # 右摇杆Y轴 - 控制RY轴旋转
            5: "RT Trigger",  # 右触发器 - Z轴向上
        }

        # 帽子开关映射（D-pad）
        self.hat_mapping = {
            (0, 1): "D-Pad Up",
            (0, -1): "D-Pad Down",
            (-1, 0): "D-Pad Left",  # 控制RZ轴负向旋转
            (1, 0): "D-Pad Right",  # 控制RZ轴正向旋转
            (-1, 1): "D-Pad Up-Left",
            (1, 1): "D-Pad Up-Right",
            (-1, -1): "D-Pad Down-Left",
            (1, -1): "D-Pad Down-Right",
        }

    def _initialize_robot(self):
        """初始化机械臂"""
        print("正在初始化机械臂...")

        try:
            # 连接机械臂
            print("步骤1: 连接机械臂...")
            self.piper.ConnectPort()
            print("机械臂连接成功")

            # 使能机械臂
            print("步骤2: 使能机械臂...")
            self.piper.EnableArm(7)
            print("机械臂使能命令已发送")

            # 使能检测
            print("步骤3: 检测使能状态...")
            self._enable_robot()

            # 移动到安全初始位姿
            print("步骤4: 移动到安全初始位姿...")
            self._move_to_initial_pose()

            # 等待移动完成并获取实际位置
            print("步骤5: 获取当前位置...")
            time.sleep(2)
            self.current_position = self._get_current_position()

            # 如果获取位置失败，使用预定义位置
            if self.current_position is None:
                print("警告: 无法获取当前位置，使用预定义位姿")
                self.current_position = self.safe_initial_pose.copy()
            else:
                print("成功获取当前位置")

            # 保存初始位置用于重置功能
            self.initial_position = self.current_position.copy()
            self.last_target_position = self.current_position.copy()
            print("初始位置已保存")

            # 显示初始位置信息
            print(
                f"初始位置: X={self.current_position[0]:.1f}, Y={self.current_position[1]:.1f}, Z={self.current_position[2]:.1f}"
            )
            print(
                f"         RX={self.current_position[3]:.1f}, RY={self.current_position[4]:.1f}, RZ={self.current_position[5]:.1f}"
            )
            print(f"V3.0超精细控制参数: 移动速度={self.move_speed}cm/s, 旋转速度={self.rotation_speed}°/s")
            print(f"                   控制频率={self.control_frequency}Hz")
            print(f"                   每帧移动距离={self.move_distance_per_frame:.3f}cm")
            print(f"                   每帧旋转角度={self.rotation_angle_per_frame:.3f}°")
            print(f"                   死区={self.deadzone}, 位置平滑因子={self.smoothness_factor}")
            print(f"                   速度平滑因子={self.velocity_smooth_factor}")
            print(f"                   非线性响应曲线: 微小控制(0.5x), 中等控制(1.0x), 大控制(1.2x)")
            print("机械臂初始化完成！")

        except Exception as e:
            print(f"机械臂初始化失败: {e}")
            import traceback

            traceback.print_exc()
            raise

    def _enable_robot(self):
        """使能机械臂并检测使能状态"""
        enable_flag = False
        timeout = 5  # 超时时间（秒）
        start_time = time.time()

        while not enable_flag:
            elapsed_time = time.time() - start_time
            print("检查使能状态...")

            try:
                # 获取所有电机的使能状态
                low_spd_info = self.piper.GetArmLowSpdInfoMsgs()
                enable_flag = (
                    low_spd_info.motor_1.foc_status.driver_enable_status
                    and low_spd_info.motor_2.foc_status.driver_enable_status
                    and low_spd_info.motor_3.foc_status.driver_enable_status
                    and low_spd_info.motor_4.foc_status.driver_enable_status
                    and low_spd_info.motor_5.foc_status.driver_enable_status
                    and low_spd_info.motor_6.foc_status.driver_enable_status
                )
            except Exception as e:
                print(f"获取使能状态失败: {e}")
                enable_flag = False

            # 如果未使能，尝试重新使能
            if not enable_flag:
                self.piper.EnableArm(7)
                self.piper.GripperCtrl(0, 1000, 0x01, 0)

            # 超时检查
            if elapsed_time > timeout:
                print("程序自动使能超时,退出程序")
                exit(0)

            time.sleep(1)
        print("机械臂使能成功")

    def _move_to_initial_pose(self):
        """移动到安全初始位姿"""
        try:
            # 将位置转换为微米单位
            X = round(self.safe_initial_pose[0] * self.factor)
            Y = round(self.safe_initial_pose[1] * self.factor)
            Z = round(self.safe_initial_pose[2] * self.factor)
            RX = round(self.safe_initial_pose[3] * self.factor)
            RY = round(self.safe_initial_pose[4] * self.factor)
            RZ = round(self.safe_initial_pose[5] * self.factor)

            print(f"移动到初始位姿: X={X}, Y={Y}, Z={Z}, RX={RX}, RY={RY}, RZ={RZ}")

            # 使用慢速移动到初始位置（速度15，确保安全）
            self.piper.MotionCtrl_2(0x01, 0x00, 15, 0x00)
            self.piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            time.sleep(3)  # 等待移动完成

        except Exception as e:
            print(f"移动到初始位姿失败: {e}")

    def _wait_motion_done(self):
        """等待机械臂运动完成"""
        start_t = time.time()
        timeout = 5  # 超时时间（秒）

        while self._is_in_motion():
            time.sleep(0.01)  # 10ms检查间隔
            if time.time() - start_t > timeout:
                print("等待运动完成超时")
                return False
        return True

    def _is_in_motion(self):
        """检查机械臂是否在运动"""
        try:
            return bool(self.piper.GetArmStatus().arm_status.motion_status)
        except:
            return False

    def _get_current_position(self):
        """获取当前机械臂末端位姿"""
        try:
            # 使用正确的API调用方法获取位姿信息
            endpose = self.piper.GetArmEndPoseMsgs().end_pose

            # 转换为毫米和度
            position = [
                endpose.X_axis / self.factor,  # 微米转毫米
                endpose.Y_axis / self.factor,
                endpose.Z_axis / self.factor,
                endpose.RX_axis / self.factor,  # 微度转度
                endpose.RY_axis / self.factor,
                endpose.RZ_axis / self.factor,
            ]
            return position

        except Exception as e:
            print(f"获取位姿失败: {e}")
            return None

    def _apply_motion_smoothing(self, target_pos):
        """
        V3.0新增：应用超精细运动平滑算法

        Args:
            target_pos: 目标位置

        Returns:
            smoothed_pos: 平滑后的位置
        """
        if not self.motion_smoothing_enabled or self.last_target_position is None:
            return target_pos

        # 计算当前速度
        current_velocity = []
        for i in range(6):
            if self.last_target_position is not None:
                vel = (target_pos[i] - self.last_target_position[i]) * self.control_frequency
            else:
                vel = 0.0
            current_velocity.append(vel)

        # 速度平滑
        if self.velocity_smoothing:
            smoothed_velocity = []
            for i in range(6):
                smoothed_vel = self.last_velocity[i] * self.velocity_smooth_factor + current_velocity[i] * (
                    1 - self.velocity_smooth_factor
                )
                smoothed_velocity.append(smoothed_vel)
            self.last_velocity = smoothed_velocity
        else:
            smoothed_velocity = current_velocity
            self.last_velocity = current_velocity

        # 位置平滑 - 使用速度积分
        smoothed_pos = []
        for i in range(6):
            current_val = self.last_target_position[i]
            # 使用平滑后的速度计算位置增量
            position_increment = smoothed_velocity[i] / self.control_frequency
            smoothed_val = current_val + position_increment
            smoothed_pos.append(smoothed_val)

        return smoothed_pos

    def _move_robot_to_position(self, new_pos, wait_completion=False):
        """移动机械臂到指定位置 - V3.0优化版本"""
        try:
            # 只有在需要等待完成时才等待（如重置位置时）
            if wait_completion:
                self._wait_motion_done()

            # V3.0新增：应用运动平滑
            if self.motion_smoothing_enabled:
                new_pos = self._apply_motion_smoothing(new_pos)

            # 转换为微米单位
            X = round(new_pos[0] * self.factor)
            Y = round(new_pos[1] * self.factor)
            Z = round(new_pos[2] * self.factor)
            RX = round(new_pos[3] * self.factor)
            RY = round(new_pos[4] * self.factor)
            RZ = round(new_pos[5] * self.factor)

            # 设置运动速度（60Hz控制频率下的连续移动）
            # arm_speed = 40  # V3.0使用更快的速度进行连续控制
            # self.piper.MotionCtrl_2(0x01, 0x00, self.arm_speed, 0x00)

            # 发送位置控制命令
            # self.piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            self.IKmove(X, Y, Z, RX, RY, RZ, speed=self.arm_speed)

            # 更新位置记录
            self.current_position = new_pos.copy()
            self.last_target_position = new_pos.copy()

        except Exception as e:
            print(f"移动失败: {e}")

    # ==================== 位置控制方法 ====================

    def move_x(self, direction):
        """X轴移动（正方向向前，负方向向后）- V3.0超精细控制"""
        if self.current_position is None:
            return

        # 应用非线性响应曲线，提升微小控制的精度
        if abs(direction) < 0.3:
            # 微小控制时使用更精细的步长
            step_multiplier = 0.5
        elif abs(direction) < 0.7:
            # 中等控制时使用标准步长
            step_multiplier = 1.0
        else:
            # 大控制时使用稍大步长
            step_multiplier = 1.2

        new_pos = self.current_position.copy()
        new_pos[0] += direction * self.move_distance_per_frame * 10 * step_multiplier  # 转换为毫米
        self._move_robot_to_position(new_pos, wait_completion=False)

    def move_y(self, direction):
        """Y轴移动（正方向向右，负方向向左）- V3.0超精细控制"""
        if self.current_position is None:
            return

        # 应用非线性响应曲线，提升微小控制的精度
        if abs(direction) < 0.3:
            step_multiplier = 0.5
        elif abs(direction) < 0.7:
            step_multiplier = 1.0
        else:
            step_multiplier = 1.2

        new_pos = self.current_position.copy()
        new_pos[1] += direction * self.move_distance_per_frame * 10 * step_multiplier  # 转换为毫米
        self._move_robot_to_position(new_pos, wait_completion=False)

    def move_z(self, direction):
        """Z轴移动（正方向向上，负方向向下）- V3.0超精细控制"""
        if self.current_position is None:
            return

        # 应用非线性响应曲线，提升微小控制的精度
        if abs(direction) < 0.3:
            step_multiplier = 0.5
        elif abs(direction) < 0.7:
            step_multiplier = 1.0
        else:
            step_multiplier = 1.2

        new_pos = self.current_position.copy()
        new_pos[2] += direction * self.move_distance_per_frame * 10 * step_multiplier  # 转换为毫米
        self._move_robot_to_position(new_pos, wait_completion=False)

    # ==================== 旋转控制方法 ====================

    def rotate_rx(self, direction):
        """RX轴旋转（绕X轴旋转）- V3.0超精细控制"""
        if self.current_position is None:
            return

        # 应用非线性响应曲线，提升微小控制的精度
        if abs(direction) < 0.3:
            step_multiplier = 0.5
        elif abs(direction) < 0.7:
            step_multiplier = 1.0
        else:
            step_multiplier = 1.2

        new_pos = self.current_position.copy()
        new_pos[3] += direction * self.rotation_angle_per_frame * step_multiplier  # 度
        self._move_robot_to_position(new_pos, wait_completion=False)

    def rotate_ry(self, direction):
        """RY轴旋转（绕Y轴旋转）- V3.0超精细控制"""
        if self.current_position is None:
            return

        # 应用非线性响应曲线，提升微小控制的精度
        if abs(direction) < 0.3:
            step_multiplier = 0.5
        elif abs(direction) < 0.7:
            step_multiplier = 1.0
        else:
            step_multiplier = 1.2

        new_pos = self.current_position.copy()
        new_pos[4] += direction * self.rotation_angle_per_frame * step_multiplier  # 度
        self._move_robot_to_position(new_pos, wait_completion=False)

    def rotate_rz(self, direction):
        """RZ轴旋转（绕Z轴旋转）- V3.0超精细控制"""
        if self.current_position is None:
            return

        # 应用非线性响应曲线，提升微小控制的精度
        if abs(direction) < 0.3:
            step_multiplier = 0.5
        elif abs(direction) < 0.7:
            step_multiplier = 1.0
        else:
            step_multiplier = 1.2

        new_pos = self.current_position.copy()
        new_pos[5] += direction * self.rotation_angle_per_frame * step_multiplier  # 度
        self._move_robot_to_position(new_pos, wait_completion=False)

    def reset_to_initial_position(self):
        """重置到初始位置"""
        print("正在重置到初始位置...")
        if self.initial_position is None:
            print("错误: 初始位置未设置，无法重置")
            return

        print(
            f"目标位置: X={self.initial_position[0]:.1f}, Y={self.initial_position[1]:.1f}, Z={self.initial_position[2]:.1f}"
        )
        print(
            f"         RX={self.initial_position[3]:.1f}, RY={self.initial_position[4]:.1f}, RZ={self.initial_position[5]:.1f}"
        )

        try:
            self._move_robot_to_position(self.initial_position, wait_completion=True)
            # self.piper.ModeCtrl(0x01, 0x01, 40, 0x00)
            # self.piper.JointCtrl(0, 0, 0, 0, 0, 0)
            # self.piper.GripperCtrl(0, 1000, 0x01, 0)

            print("已重置到初始位置")
        except Exception as e:
            print(f"重置位置失败: {e}")
            import traceback

            traceback.print_exc()

    # ==================== 夹爪控制方法 ====================

    def open_gripper(self):
        """打开夹爪（100mm宽度）"""
        print("正在打开夹爪...")
        try:
            # 等待运动完成
            self._wait_motion_done()
            # 发送夹爪控制命令 - 100mm宽度
            self.piper.GripperCtrl(100 * 1000, 1000, 0x01, 0)
            # 等待夹爪操作完成
            time.sleep(self.gripper_delay)
            print("夹爪已打开")
        except Exception as e:
            print(f"打开夹爪失败: {e}")

    def close_gripper(self):
        """关闭夹爪（完全闭合）"""
        print("正在关闭夹爪...")
        try:
            # 等待运动完成
            self._wait_motion_done()
            # 发送夹爪控制命令 - 0mm宽度（完全闭合）
            self.piper.GripperCtrl(0, 1000, 0x01, 0)
            # 等待夹爪操作完成
            time.sleep(self.gripper_delay)
            print("夹爪已关闭")
        except Exception as e:
            print(f"关闭夹爪失败: {e}")

    # ==================== Xbox手柄控制方法 ====================

    def handle_joystick_control(self, axes, hats):
        """
        处理摇杆和D-pad控制 - V3.0优化版本

        Args:
            axes: 手柄轴数据列表
            hats: 手柄帽子数据列表
        """
        if len(axes) < 6:
            return

        # 获取摇杆和触发器值
        left_stick_x = -axes[0]  # 左摇杆X轴（取反修正左右方向）
        left_stick_y = -axes[1]  # 左摇杆Y轴（反向）
        lt_trigger = axes[2]  # 左触发器
        right_stick_x = -axes[3]  # 右摇杆X轴（取反修正方向）
        right_stick_y = axes[4]  # 右摇杆Y轴（修正方向）
        rt_trigger = axes[5]  # 右触发器

        # 获取D-pad状态
        dpad_x = 0
        dpad_y = 0
        if len(hats) > 0:
            dpad_x, dpad_y = hats[0]

        # 应用死区 - V3.0优化死区处理
        if abs(left_stick_x) < self.deadzone:
            left_stick_x = 0
        if abs(left_stick_y) < self.deadzone:
            left_stick_y = 0
        if abs(right_stick_x) < self.deadzone:
            right_stick_x = 0
        if abs(right_stick_y) < self.deadzone:
            right_stick_y = 0

        # 触发器特殊处理 - Xbox手柄触发器未按下时为-1，按下时为1
        if lt_trigger < 0:
            lt_trigger = 0  # 未按下状态设为0
        elif abs(lt_trigger) < self.deadzone:
            lt_trigger = 0  # 死区内设为0

        if rt_trigger < 0:
            rt_trigger = 0  # 未按下状态设为0
        elif abs(rt_trigger) < self.deadzone:
            rt_trigger = 0  # 死区内设为0

        # ==================== 位置控制 ====================
        # X轴移动控制（左摇杆Y轴）
        if abs(left_stick_y) > self.deadzone:
            self.move_x(left_stick_y)

        # Y轴移动控制（左摇杆X轴）
        if abs(left_stick_x) > self.deadzone:
            self.move_y(left_stick_x)

        # Z轴移动控制（LT/RT触发器）
        if lt_trigger > 0:
            # LT触发器 - Z轴向下
            self.move_z(-lt_trigger)
        elif rt_trigger > 0:
            # RT触发器 - Z轴向上
            self.move_z(rt_trigger)

        # ==================== 旋转控制 ====================
        # RZ轴旋转控制（D-pad Left/Right）
        if dpad_x != 0:
            self.rotate_rz(dpad_x)

        # RZ轴旋转控制（右摇杆X轴）
        if abs(right_stick_x) > self.deadzone:
            self.rotate_rz(right_stick_x)

        # RY轴旋转控制（右摇杆Y轴）
        if abs(right_stick_y) > self.deadzone:
            self.rotate_ry(right_stick_y)

        # V3.0新增：超详细的控制状态反馈
        active_controls = []
        control_intensity = []

        if abs(left_stick_x) > 0.01 or abs(left_stick_y) > 0.01:
            active_controls.append(f"位置: X={left_stick_y:.3f}, Y={left_stick_x:.3f}")
            # 计算控制强度
            pos_intensity = max(abs(left_stick_x), abs(left_stick_y))
            if pos_intensity < 0.3:
                control_intensity.append("精细")
            elif pos_intensity < 0.7:
                control_intensity.append("中等")
            else:
                control_intensity.append("快速")

        if lt_trigger > 0.01 or rt_trigger > 0.01:
            z_control = rt_trigger - lt_trigger
            active_controls.append(f"Z={z_control:.3f}")
            if abs(z_control) < 0.3:
                control_intensity.append("精细")
            elif abs(z_control) < 0.7:
                control_intensity.append("中等")
            else:
                control_intensity.append("快速")

        if dpad_x != 0:
            active_controls.append(f"RZ(D-pad)={dpad_x}")
            control_intensity.append("步进")

        if abs(right_stick_x) > 0.01 or abs(right_stick_y) > 0.01:
            active_controls.append(f"旋转: RZ(右X)={right_stick_x:.3f}, RY(右Y)={right_stick_y:.3f}")
            # 计算旋转控制强度
            rot_intensity = max(abs(right_stick_x), abs(right_stick_y))
            if rot_intensity < 0.3:
                control_intensity.append("精细")
            elif rot_intensity < 0.7:
                control_intensity.append("中等")
            else:
                control_intensity.append("快速")

        if active_controls:
            intensity_str = " | ".join(control_intensity) if control_intensity else "待机"
            print(f"V3.0控制状态: {' | '.join(active_controls)} | 模式: {intensity_str}")

    def handle_button_control(self, buttons):
        """
        处理按钮控制 - 夹爪和重置

        Args:
            buttons: 手柄按钮状态列表
        """
        if len(buttons) < 8:
            current_buttons = buttons + [False] * (8 - len(buttons))
        else:
            current_buttons = buttons[:8]

        # 检查按钮状态变化
        for i, button_pressed in enumerate(current_buttons):
            last_pressed = self.last_button_states.get(i, False)

            if button_pressed and not last_pressed:  # 按钮被按下
                print(f"按钮 {i} 被按下")
                if i == 4:  # LB按钮 - 关闭夹爪
                    print("LB按钮 - 关闭夹爪")
                    self.close_gripper()
                elif i == 5:  # RB按钮 - 打开夹爪
                    print("RB按钮 - 打开夹爪")
                    self.open_gripper()
                elif i == 7:  # Menu按钮 - 重置位置
                    print("Menu按钮 - 重置位置")
                    self.reset_to_initial_position()

            # 更新按钮状态
            self.last_button_states[i] = button_pressed

    def update_performance_monitor(self):
        """V3.0新增：性能监控"""
        self.frame_count += 1
        current_time = time.time()

        if current_time - self.last_fps_check >= 1.0:  # 每秒更新一次
            self.actual_fps = self.frame_count / (current_time - self.last_fps_check)
            self.frame_count = 0
            self.last_fps_check = current_time

    def read_xbox_inputs(self):
        """读取Xbox手柄输入 - V3.0优化版本"""
        # 处理pygame事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False

        # 读取摇杆、按钮和帽子状态
        axes = [self.joystick.get_axis(i) for i in range(self.joystick.get_numaxes())]
        buttons = [self.joystick.get_button(i) for i in range(self.joystick.get_numbuttons())]
        hats = [self.joystick.get_hat(i) for i in range(self.joystick.get_numhats())]

        # 处理摇杆和D-pad控制
        self.handle_joystick_control(axes, hats)

        # 处理按钮控制
        self.handle_button_control(buttons)

        # V3.0新增：性能监控
        self.update_performance_monitor()

        return True

    def run(self):
        """运行主控制循环 - V3.0优化版本"""
        print("\n=============== Piper机械臂Xbox控制系统 V3.0 ===============")
        print("作者: Yinchuan Li")
        print("日期: 2025-07-27")
        print("\nV3.0超精细版本升级说明:")
        print("  - 控制频率保持60Hz，确保响应速度")
        print("  - 移动速度进一步降低到2cm/s，控制更加精细")
        print("  - 旋转速度进一步降低到10°/s，旋转更加平滑")
        print("  - 死区进一步降低到0.03，提升微小控制响应性")
        print("  - 新增速度平滑算法，减少机械臂抖动")
        print("  - 非线性响应曲线，提升控制精度")
        print("  - 完整的6自由度控制能力")
        print("  - 超精细控制精度和丝滑运动体验")
        print("\n控制说明:")
        print("  [位置控制 - 持续移动]")
        print("  - 左摇杆X轴: 控制机械臂Y轴左右移动")
        print("  - 左摇杆Y轴: 控制机械臂X轴前后移动")
        print("  - LT触发器: Z轴向下移动")
        print("  - RT触发器: Z轴向上移动")
        print("  [旋转控制 - 持续旋转]")
        print("  - D-pad Left/Right: 控制RZ轴旋转")
        print("  - 右摇杆X轴: 控制RZ轴旋转")
        print("  - 右摇杆Y轴: 控制RY轴旋转")
        print("  [夹爪控制]")
        print("  - LB按钮: 关闭夹爪")
        print("  - RB按钮: 打开夹爪")
        print("  - Menu按钮: 重置到初始位置")
        print("  [退出控制]")
        print("  - ESC键: 退出程序")
        print("\nV3.0超精细优化特性:")
        print("  - 60Hz控制频率，确保响应速度")
        print("  - 2cm/s移动速度，10°/s旋转速度")
        print("  - 速度平滑算法，减少抖动")
        print("  - 非线性响应曲线，提升控制精度")
        print("  - 超精细控制，微小控制步长减小50%")
        print("  - 实时性能监控和状态显示")
        print("=" * 65)

        try:
            # 主控制循环 - 60Hz频率
            control_interval = 1.0 / self.control_frequency  # 约16.67ms
            last_control_time = time.time()

            while True:
                current_time = time.time()

                # 确保60Hz控制频率
                if current_time - last_control_time >= control_interval:
                    # 读取手柄输入
                    if not self.read_xbox_inputs():
                        break

                    # 显示实时状态信息
                    if self.current_position is not None:
                        print(
                            f"\r当前位置: X={self.current_position[0]:6.1f}, Y={self.current_position[1]:6.1f}, Z={self.current_position[2]:6.1f} | "
                            f"RX={self.current_position[3]:6.1f}, RY={self.current_position[4]:6.1f}, RZ={self.current_position[5]:6.1f} | "
                            f"FPS: {self.actual_fps:4.1f}Hz",
                            end="",
                            flush=True,
                        )

                    last_control_time = current_time
                else:
                    # 短暂休眠以节省CPU
                    time.sleep(0.001)  # 1ms

        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
        except Exception as e:
            print(f"\n\n程序运行出错: {e}")
            import traceback

            traceback.print_exc()
        finally:
            # 清理资源
            print("\n正在清理资源...")
            try:
                pygame.quit()
                print("Xbox手柄已断开连接")
            except:
                pass
            print("程序已退出")

class XboxPiperControllerV4:
    """
    V4.0: 支持一个手柄切换控制两台Piper机械臂（can0/can1），A/B键切换
    其余控制逻辑与V3.0一致
    """

    def __init__(self, piper0, piper1):
        """
        初始化，接收两台机械臂接口
        """
        self.piper0 = piper0
        self.piper1 = piper1
        self.active_piper = self.piper0  # 默认控制can0
        self.active_piper_name = "can0"
        self.controllers = {
            "can0": XboxPiperControllerV3(self.piper0),
            "can1": XboxPiperControllerV3(self.piper1),
        }
        self.active_controller = self.controllers[self.active_piper_name]

        # 初始化手柄
        self._init_xbox_controller()

    def _init_xbox_controller(self):
        pygame.init()
        pygame.joystick.init()
        if pygame.joystick.get_count() == 0:
            raise RuntimeError("未检测到Xbox手柄，请确保手柄已连接")
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        print(f"已连接Xbox手柄: {self.joystick.get_name()}")

    def switch_arm(self, arm_name):
        if arm_name not in self.controllers:
            print(f"未知机械臂: {arm_name}")
            return
        if self.active_piper_name != arm_name:
            self.active_piper_name = arm_name
            self.active_controller = self.controllers[arm_name]
            print(f"\n>>> 已切换为控制 {arm_name} 机械臂 <<<\n")

    def run(self):
        print("\n=============== Piper机械臂Xbox控制系统 V4.0 ===============")
        print("A键：切换为控制can0机械臂")
        print("B键：切换为控制can1机械臂")
        print("其余控制与V3.0一致")
        print("=" * 65)
        try:
            control_interval = 1.0 / self.active_controller.control_frequency
            last_control_time = time.time()
            while True:
                current_time = time.time()
                if current_time - last_control_time >= control_interval:
                    # 读取手柄输入
                    for event in pygame.event.get():
                        if event.type == pygame.QUIT:
                            return
                        elif event.type == pygame.KEYDOWN:
                            if event.key == pygame.K_ESCAPE:
                                return
                    # 读取按钮
                    buttons = [self.joystick.get_button(i) for i in range(self.joystick.get_numbuttons())]
                    # 检查A/B键切换
                    if len(buttons) >= 2:
                        if buttons[0]:  # A键
                            self.switch_arm("can0")
                        elif buttons[1]:  # B键
                            self.switch_arm("can1")
                    # 读取摇杆、帽子
                    axes = [self.joystick.get_axis(i) for i in range(self.joystick.get_numaxes())]
                    hats = [self.joystick.get_hat(i) for i in range(self.joystick.get_numhats())]
                    # 交给当前控制器处理
                    self.active_controller.handle_joystick_control(axes, hats)
                    self.active_controller.handle_button_control(buttons)
                    self.active_controller.update_performance_monitor()
                    # 状态显示
                    pos = self.active_controller.current_position
                    if pos is not None:
                        print(
                            f"\r[{self.active_piper_name}] X={pos[0]:6.1f}, Y={pos[1]:6.1f}, Z={pos[2]:6.1f} | "
                            f"RX={pos[3]:6.1f}, RY={pos[4]:6.1f}, RZ={pos[5]:6.1f} | "
                            f"FPS: {self.active_controller.actual_fps:4.1f}Hz",
                            end="", flush=True
                        )
                    last_control_time = current_time
                else:
                    time.sleep(0.001)
        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
        except Exception as e:
            print(f"\n\n程序运行出错: {e}")
            import traceback
            traceback.print_exc()
        finally:
            print("\n正在清理资源...")
            try:
                pygame.quit()
                print("Xbox手柄已断开连接")
            except:
                pass
            print("程序已退出")

# 其余类（XboxPiperControllerV3、MujocoIKSolver等）可直接复用V3.0代码，无需修改

def main():
    """主函数"""
    try:
        # 创建两台Piper机械臂对象
        piper0 = C_PiperInterface("can0")
        piper1 = C_PiperInterface("can1")
        # 创建V4控制器
        controller = XboxPiperControllerV4(piper0, piper1)
        controller.run()
    except Exception as e:
        print(f"程序启动失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
# ...existing code...